import geatpy as ea
import numpy as np
import random
from math import floor


# 生成初始种群，不允许有重复的基因
def generate(pop_num, chrom_len, lb, ub):
    # 生成[lb, ub]之间的所有整数
    index = np.arange(lb, ub+1, 1, dtype=int)
    pop = np.zeros((pop_num, chrom_len), dtype=int)
    for i in range(pop_num):
        pop[i, :] = np.random.choice(index, size=chrom_len, replace=False)
    return pop


# 适应度函数，score越小，适应度越大
# score必须是一个n行1列的向量
def rank(score):
    return ea.ranking(score)


# stochastic universal sampling
# fitness越大越好
def select(fitness, sel_num):
    """
    随机遍历选择
    :param fitness: 形状为n*1的numpy array，存储着对应个体的适应度。越大越好
    :param sel_num: 整型，被选择出来的数量
    :return: 返回被选择个体的下标
    """
    return ea.sus(fitness, sel_num, None)


def uniform_crossover(pop, pb):
    """
    均匀交叉
    :param pop: n行m列的numpy的array，待交叉的种群
    :param pb: 浮点型，交叉的概率
    :return: n行m列的numpy的array，交叉后的种群
    """
    # 种群数量
    pop_num = pop.shape[0]
    # 染色体长度
    chrom_len = pop.shape[1]
    # 交叉的两个染色体
    pair = match(pop_num)
    # 对于每对需要交叉的染色体
    for elem in pair:
        # prob为交叉概率
        if random.random() < pb:
            # chrom1和chrom2是numpy的数组，形状为(n, )，即一维数组
            # 表示待交叉的两个染色体
            chrom1 = pop[elem[0], :]
            chrom2 = pop[elem[1], :]
            # 染色体中的每个基因
            for i in range(chrom_len):
                if (chrom1[i] not in chrom2) and (chrom2[i] not in chrom1):
                    # 每个基因发生交叉的概率为0.5
                    if random.random() < 0.5:
                        pop[elem[0], i], pop[elem[1], i] = pop[elem[1], i], pop[elem[0], i]
    return pop


# 暂时有问题，还没有去重
def two_point_crossover(pop, pb):
    """
    两点交叉
    :param pop: n行m列的numpy的array，待交叉的种群
    :param pb: 浮点型，交叉的概率
    :return: n行m列的numpy的array，交叉后的种群
    """
    # 种群数量
    pop_num = pop.shape[0]
    # 染色体长度
    chrom_len = pop.shape[1]
    # 交叉的两个染色体
    pairs = match(pop_num)
    # 对于每对需要交叉的染色体
    for pair in pairs:
        # prob为交叉概率
        if random.random() < pb:
            # 两个交叉点
            point1 = random.randint(1, chrom_len)
            point2 = random.randint(1, chrom_len-1)
            # 保持point1小于point2
            if point2 >= point1:
                point2 += 1
            else:
                point1, point2 = point2, point1
            # 交换[point1, point2-1]之间的染色体片段
            pop[pair[0], point1:point2], pop[pair[0], point1:point2] \
                = pop[pair[0], point1:point2],  pop[pair[0], point1:point2]
    return pop


def mutate(pop, lb, ub, pb):
    """
    均匀变异
    :param pop: 待变异种群， n行m列的numpy的array
    :param lb: 下界，整型
    :param ub: 上界，整型
    :param pb: 变异概率
    :return: 变异后的种群，n行m列的numpy的array
    """
    pop_num = pop.shape[0]
    chrom_len = pop.shape[1]
    low = tuple(lb for i in range(chrom_len))
    high = tuple(ub for i in range(chrom_len))
    index = tuple(range(chrom_len))
    for r in range(pop_num):
        for c, xl, xu in zip(index, low, high):
            if random.random() < pb:
                pop[r, c] = unique_choice(pop[r, :], lb, ub)
    return pop


def mutate_gaussian():
    pass



def unique_choice(chrom, lb, ub):
    """

    :param chrom: numpy的array，形状为(n, )，染色体数组
    :param lb: 整型，下界
    :param ub: 整型，下界
    :return: 返回一个不再chrom中而再[lb, ub]中的随机整数
    """
    index_set = set(i for i in range(lb, ub+1))
    for gene in chrom:
        index_set.remove(gene)
    return random.choice(tuple(index_set))
    # in_pop = dict()
    # for i in range(lb, ub+1):
    #     in_pop[i] = False
    # for r in range(pop.shape[0]):
    #     for c in range(pop.shape[1]):
    #         in_pop[pop[r, c]] = True
    # selecting_set = []
    # for item in in_pop.items():
    #     if item[1] == False:
    #         selecting_set.append(item[0])
    # print("ERRRRRRRRRRRRRROR: ", selecting_set)
    # return np.random.choice(selecting_set, 1)



def match(pop_num):
    index_set = set(i for i in range(pop_num))
    result = []
    for i in range(floor(pop_num/2)):
        num1 = random.choice(tuple(index_set))
        index_set.remove(num1)
        num2 = random.choice(tuple(index_set))
        index_set.remove(num2)
        result.append((num1, num2))
    return tuple(result)



if __name__ == "__main__":
    pop_num = 5
    chrom_len = 5
    lb = 0
    ub = 9
    iter_num = 10
    pop = generate(pop_num, chrom_len, lb, ub)
    print("初始种群：")
    print(pop)
    for _ in range(iter_num):
        print("---------------------------------------------")
        score = np.random.random(pop_num)
        score = score.reshape(pop_num, 1)
        fitness = rank(score)
        sel_index = select(fitness, pop_num)
        pop = pop[sel_index, :]
        print("选择：")
        print(pop)
        pop = two_point_crossover(pop, 0.9)
        print("交叉：")
        print(pop)
        pop = mutate(pop, lb, ub, 1/chrom_len)
        print("变异：")
        print(pop)
